import cv2
from concurrent.futures import ThreadPoolExecutor
import os
from tqdm import tqdm

class VideoProcessor:
    def __init__(self, video_path, output_folder, rotate=0, batch_size=50):
        self.video_path = video_path
        self.output_folder = output_folder
        self.rotate = rotate
        self.batch_size = batch_size
        self.first_frame_saved = False  # 添加标志位
        
    def read_frames(self):
        """读取视频帧"""
        cap = cv2.VideoCapture(self.video_path)
        if not cap.isOpened():
            raise ValueError(f"无法打开视频文件: {self.video_path}")
        
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_skip = max(1, int(fps / 10))
        print(f"\n开始读取视频帧... (总帧数: {total_frames}, FPS: {fps:.2f}) frame_skip: {frame_skip}")
        frames = []
        pbar = tqdm(total=total_frames, desc="读取进度")
        
        frame_count = 0
        while frame_count < total_frames:
            ret, frame = cap.read()
            if not ret:
                break
                
            if frame_count % frame_skip == 0:
                if self.rotate:
                    frame = self._rotate_frame(frame, self.rotate, is_first=(not self.first_frame_saved))
                    self.first_frame_saved = True
                frames.append((frame_count, frame))
                
                if len(frames) >= self.batch_size:
                    yield frames
                    frames = []
            
            frame_count += 1
            pbar.update(1)
                
        if frames:
            yield frames
            
        cap.release()
        pbar.close()
    
    def _rotate_frame(self, frame, angle, is_first=False):
        """旋转帧"""
        rotated_frame = frame
        if angle == 90:
            rotated_frame = cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE)
        elif angle == 180:
            rotated_frame = cv2.rotate(frame, cv2.ROTATE_180)
        elif angle == 270:
            rotated_frame = cv2.rotate(frame, cv2.ROTATE_90_COUNTERCLOCKWISE)
        
        # 保存第一帧用于调试
        if is_first:
            debug_dir = os.path.join(self.output_folder, 'debug')
            os.makedirs(debug_dir, exist_ok=True)
            
            # 同时保存原始帧和旋转后的帧
            orig_debug_path = os.path.join(debug_dir, 'first_frame_original.jpg')
            rot_debug_path = os.path.join(debug_dir, f'first_frame_rotated_{angle}.jpg')
            
            cv2.imwrite(orig_debug_path, frame)
            cv2.imwrite(rot_debug_path, rotated_frame)
            print(f"\n已保存原始第一帧到: {orig_debug_path}")
            print(f"已保存旋转后的第一帧到: {rot_debug_path}")
            
        return rotated_frame